{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting MNIST_data\\train-images-idx3-ubyte.gz\n",
      "Extracting MNIST_data\\train-labels-idx1-ubyte.gz\n",
      "Extracting MNIST_data\\t10k-images-idx3-ubyte.gz\n",
      "Extracting MNIST_data\\t10k-labels-idx1-ubyte.gz\n",
      "157 / 859: -0.194456, 1.057701"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "%matplotlib inline \n",
    "import sys\n",
    "from tensorflow.examples.tutorials.mnist import input_data\n",
    "mnist = input_data.read_data_sets('MNIST_data', one_hot=True)\n",
    "\n",
    "def plot(samples, D_loss, G_loss, epoch, total):\n",
    "\n",
    "    fig = plt.figure(figsize=(10, 5))        \n",
    "\n",
    "    gs = gridspec.GridSpec(4, 8)\n",
    "    gs.update(wspace=0.05, hspace=0.05)\n",
    "    \n",
    "    # Plot losses\n",
    "    ax = plt.subplot(gs[:, 4:])\n",
    "    ax.plot(D_loss, label=\"discriminator's loss\", color='b')\n",
    "    ax.plot(G_loss, label=\"generator's loss\", color='r')\n",
    "    ax.set_xlim([0, total])\n",
    "    ax.yaxis.tick_right()\n",
    "    ax.legend()\n",
    "\n",
    "    # Generate images\n",
    "    for i, sample in enumerate(samples):\n",
    "\n",
    "        if i > 4* 4 - 1:\n",
    "            break\n",
    "        ax = plt.subplot(gs[i % 4, int(i / 4)])\n",
    "        plt.axis('off')\n",
    "        ax.set_xticklabels([])\n",
    "        ax.set_yticklabels([])\n",
    "        ax.set_aspect('equal')\n",
    "        plt.imshow(sample.reshape(28, 28), cmap='Greys_r')\n",
    "\n",
    "    plt.savefig('./output/' + str(epoch + 1) + '.png')\n",
    "    plt.close()\n",
    "In [4]:\n",
    "def Conv2d(input, output_dim=64, kernel=(5, 5), strides=(2, 2), stddev=0.2, name='conv_2d'):\n",
    "\n",
    "    with tf.variable_scope(name):\n",
    "        W = tf.get_variable('Conv2dW', [kernel[0], kernel[1], input.get_shape()[-1], output_dim],\n",
    "                             initializer=tf.truncated_normal_initializer(stddev=stddev))\n",
    "        b = tf.get_variable('Conv2db', [output_dim], initializer=tf.zeros_initializer())\n",
    "        \n",
    "        return tf.nn.conv2d(input, W, strides=[1, strides[0], strides[1], 1], padding='SAME') + b\n",
    "\n",
    "def Deconv2d(input, output_dim, batch_size, kernel=(5, 5), strides=(2, 2), stddev=0.2, name='deconv_2d'):\n",
    "    \n",
    "    with tf.variable_scope(name):\n",
    "        W = tf.get_variable('Deconv2dW', [kernel[0], kernel[1], output_dim, input.get_shape()[-1]],\n",
    "                             initializer=tf.truncated_normal_initializer(stddev=stddev))\n",
    "        b = tf.get_variable('Deconv2db', [output_dim], initializer=tf.zeros_initializer())\n",
    "\n",
    "        input_shape = input.get_shape().as_list()\n",
    "        output_shape = [batch_size,\n",
    "                        int(input_shape[1] * strides[0]),\n",
    "                        int(input_shape[2] * strides[1]),\n",
    "                        output_dim]\n",
    "\n",
    "        deconv = tf.nn.conv2d_transpose(input, W, output_shape=output_shape,\n",
    "                                        strides=[1, strides[0], strides[1], 1])\n",
    "    \n",
    "        return deconv + b\n",
    "\n",
    "def Dense(input, output_dim, stddev=0.02, name='dense'):\n",
    "    \n",
    "    with tf.variable_scope(name):\n",
    "        \n",
    "        shape = input.get_shape()\n",
    "        W = tf.get_variable('DenseW', [shape[1], output_dim], tf.float32,\n",
    "                            tf.random_normal_initializer(stddev=stddev))\n",
    "        b = tf.get_variable('Denseb', [output_dim],\n",
    "                            initializer=tf.zeros_initializer())\n",
    "        \n",
    "        return tf.matmul(input, W) + b\n",
    "\n",
    "def BatchNormalization(input, name='bn'):\n",
    "    \n",
    "    with tf.variable_scope(name):\n",
    "    \n",
    "        output_dim = input.get_shape()[-1]\n",
    "        beta = tf.get_variable('BnBeta', [output_dim],\n",
    "                            initializer=tf.zeros_initializer())\n",
    "        gamma = tf.get_variable('BnGamma', [output_dim],\n",
    "                            initializer=tf.ones_initializer())\n",
    "    \n",
    "        if len(input.get_shape()) == 2:\n",
    "            mean, var = tf.nn.moments(input, [0])\n",
    "        else:\n",
    "            mean, var = tf.nn.moments(input, [0, 1, 2])\n",
    "        return tf.nn.batch_normalization(input, mean, var, beta, gamma, 1e-5)\n",
    "    \n",
    "def LeakyReLU(input, leak=0.2, name='lrelu'):\n",
    "    \n",
    "    return tf.maximum(input, leak*input)\n",
    "In [5]:\n",
    "BATCH_SIZE = 64\n",
    "EPOCHS = 40\n",
    "\n",
    "def Discriminator(X, reuse=False, name='d'):\n",
    "\n",
    "    with tf.variable_scope(name, reuse=reuse):\n",
    "\n",
    "        if len(X.get_shape()) > 2:\n",
    "            # X: -1, 28, 28, 1\n",
    "            D_conv1 = Conv2d(X, output_dim=64, name='D_conv1')\n",
    "        else:\n",
    "            D_reshaped = tf.reshape(X, [-1, 28, 28, 1])\n",
    "            D_conv1 = Conv2d(D_reshaped, output_dim=64, name='D_conv1')\n",
    "        D_h1 = LeakyReLU(D_conv1) # [-1, 28, 28, 64]\n",
    "        D_conv2 = Conv2d(D_h1, output_dim=128, name='D_conv2')\n",
    "        D_h2 = LeakyReLU(D_conv2) # [-1, 28, 28, 128]\n",
    "        D_r2 = tf.reshape(D_h2, [-1, 256])\n",
    "        D_h3 = LeakyReLU(D_r2) # [-1, 256]\n",
    "        D_h4 = tf.nn.dropout(D_h3, 0.5)\n",
    "        D_h5 = Dense(D_h4, output_dim=1, name='D_h5') # [-1, 1]\n",
    "        return tf.nn.sigmoid(D_h5)\n",
    "\n",
    "def Generator(z, name='g'):\n",
    "\n",
    "    with tf.variable_scope(name):\n",
    "\n",
    "        G_1 = Dense(z, output_dim=1024, name='G_1') # [-1, 1024]\n",
    "        G_bn1 = BatchNormalization(G_1, name='G_bn1')\n",
    "        G_h1 = tf.nn.relu(G_bn1)\n",
    "        G_2 = Dense(G_h1, output_dim=7*7*128, name='G_2') # [-1, 7*7*128]\n",
    "        G_bn2 = BatchNormalization(G_2, name='G_bn2')        \n",
    "        G_h2 = tf.nn.relu(G_bn2)\n",
    "        G_r2 = tf.reshape(G_h2, [-1, 7, 7, 128])\n",
    "        G_conv3 = Deconv2d(G_r2, output_dim=64, batch_size=BATCH_SIZE, name='G_conv3')\n",
    "        G_bn3 = BatchNormalization(G_conv3, name='G_bn3')        \n",
    "        G_h3 = tf.nn.relu(G_bn3)\n",
    "        G_conv4 = Deconv2d(G_h3, output_dim=1, batch_size=BATCH_SIZE, name='G_conv4')\n",
    "        G_r4 = tf.reshape(G_conv4, [-1, 784])\n",
    "        return tf.nn.sigmoid(G_r4)\n",
    "\n",
    "X = tf.placeholder(tf.float32, shape=[None, 784])\n",
    "z = tf.placeholder(tf.float32, shape=[None, 100])\n",
    "\n",
    "G = Generator(z, 'G')\n",
    "D_real = Discriminator(X, False, 'D')\n",
    "D_fake = Discriminator(G, True, 'D')\n",
    "\n",
    "D_loss = -tf.reduce_mean(tf.log(D_real) - tf.log(D_fake)) # Train to judge if the data is real correctly\n",
    "G_loss = -tf.reduce_mean(tf.log(D_fake)) # Train to pass the discriminator as real data\n",
    "\n",
    "vars = tf.trainable_variables()\n",
    "d_params = [v for v in vars if v.name.startswith('D/')]\n",
    "g_params = [v for v in vars if v.name.startswith('G/')]\n",
    "\n",
    "D_solver = tf.train.AdamOptimizer(learning_rate=1e-4, beta1=0.1).minimize(D_loss, var_list=d_params)\n",
    "G_solver = tf.train.AdamOptimizer(learning_rate=2e-4, beta1=0.3).minimize(G_loss, var_list=g_params)\n",
    "\n",
    "with tf.Session() as sess:\n",
    "\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    \n",
    "    D_loss_vals = []\n",
    "    G_loss_vals = []\n",
    "\n",
    "    iteration = int(mnist.train.images.shape[0] / BATCH_SIZE)\n",
    "    for e in range(EPOCHS):\n",
    "\n",
    "        for i in range(iteration):\n",
    "            x, _ = mnist.train.next_batch(BATCH_SIZE)\n",
    "            rand = np.random.uniform(0., 1., size=[BATCH_SIZE, 100])\n",
    "            _, D_loss_curr = sess.run([D_solver, D_loss], {X: x, z: rand})\n",
    "            rand = np.random.uniform(0., 1., size=[BATCH_SIZE, 100])\n",
    "            _, G_loss_curr = sess.run([G_solver, G_loss], {z: rand})\n",
    "\n",
    "            D_loss_vals.append(D_loss_curr)\n",
    "            G_loss_vals.append(G_loss_curr)\n",
    "\n",
    "            sys.stdout.write(\"\\r%d / %d: %f, %f\" % (i, iteration, D_loss_curr, G_loss_curr))\n",
    "            sys.stdout.flush()\n",
    "\n",
    "        data = sess.run(G, {z: rand})\n",
    "        plot(data, D_loss_vals, G_loss_vals, e, EPOCHS * iteration)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
